package com.javatest.springboot2ratelimit.redis;

import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisURI;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.dynamic.Commands;
import io.lettuce.core.dynamic.RedisCommandFactory;
import io.lettuce.core.dynamic.annotation.Command;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author lzy
 * @create 2020 -07-22 上午11:05
 */
public class RedisCellRateLimitUtil {
    /**
     * The constant GET_TOKEN_SUCCESS.
     */
    private static final String GET_TOKEN_SUCCESS = "0";
    /**
     * The constant commands.
     */
    private static final RateLimitCommand commands;
    /**
     * max_burst初始大小
     */
    private static final long initCapacity = 20;
    /**
     * tokens 操作数上限
     */
    private static final long limitCount = 20;
    /**
     * seconds 限流时长,单位秒
     */
    private static final long limitPeriod = 1;
    /**
     * apply 1 token, 每次取token的数目,默认为1
     */
    private final static long getCount = 1;

    static {
        // 创建单机连接的连接信息
        RedisURI redisUri = RedisURI.builder()
                .withHost("localhost")
                .withPort(6379)
                .withTimeout(Duration.of(10, ChronoUnit.SECONDS))
                .build();
        // 创建连接（会自动重连）
        final RedisClient redisClient = RedisClient.create(redisUri);
        final StatefulRedisConnection<String, String> conn = redisClient.connect();
        // 获取Command
        RedisCommandFactory factory = new RedisCommandFactory(conn);
        commands = factory.getCommands(RateLimitCommand.class);
    }

    /**
     * Is allow operate boolean.
     *
     * @param name the name
     * @return the boolean
     */
    public static boolean isAllowOperate(String name) {
        //最多等两次
        for (int i = 0; i < 3; i++) {
            List<Object> resultList = commands.throttle(name, initCapacity, limitCount, limitPeriod, getCount);
            if (GET_TOKEN_SUCCESS.equals(resultList.get(0).toString())) {
                return true;
            }

            //睡半秒
            try {
                TimeUnit.MILLISECONDS.sleep(500);
            } catch (InterruptedException e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
            }
        }

        System.out.println("fail," + name);
        return false;
    }

    /**
     * Is allow operate boolean.
     *
     * @param name        the name
     * @param limitCount  the operation count
     * @param limitPeriod the second count
     * @return the boolean
     */
    public static boolean isAllowOperate(String name, long limitCount, long limitPeriod) {
        long initCount = limitCount / 2;
        List<Object> resultList = commands.throttle(name, initCount, limitCount, limitPeriod, getCount);
        return GET_TOKEN_SUCCESS.equals(resultList.get(0).toString());
    }

    /**
     * 自定义封装Redis-cell提供的cl.throttle指令
     */
    interface RateLimitCommand extends Commands {
        /**
         *  限流，如CL.THROTTLE test 10 5 60 1
         *  test是key名字，初始10个令牌，每60秒允许使用5个令牌，本次获取一个令牌
         *  127.0.0.1:6379> CL.THROTTLE test 10 5 60 1
         *  1) (integer) 0 0成功，1拒绝
         *  2) (integer) 11 令牌桶的容量
         *  3) (integer) 10 当前令牌桶中可用的令牌
         *  4) (integer) -1 若请求被拒绝，这个值表示多久后才令牌桶中会重新添加令牌，单位：秒，可以作为重试时间
         *  5) (integer) 12 表示多久后令牌桶中的令牌会存满
         *
         * @param key            the key
         * @param initCapacity   the init capacity
         * @param operationCount the operation count
         * @param secondCount    the second count
         * @param getCount       the get count
         * @return the list
         */
        @Command("cl.throttle ?0 ?1 ?2 ?3 ?4")
        List<Object> throttle(String key, long initCapacity, long operationCount, long secondCount, long getCount);
    }
}
